# -*- coding: utf-8 -*-
# __author__: yb
# Version 1.0
# Contact: yb890102@126.com

import random, binascii, time
import socket, struct

class ip_pkt():
    #IP header
    def __init__(self, sip, dip):
        self.ver_hlg = 0x45
        self.DSCP = 0x00
        self.identification = 0x292a
        self.flagment0ffset = 0x4000
        self.TTL = 64
        self.protocol = 0x11
        self.headCheckSum = 0x0000
        self.sourceBuff = socket.inet_aton(sip)
        self.destinationBuff = socket.inet_aton(dip)
        self.data = ''
    def get_pkt_length(self):
        return 20 + len(self.data)
    def _temp_pack(self):
        #not insert head checksum
        tempbuff = struct.pack('!BBHHHBBH4s4s', self.ver_hlg, self.DSCP, self.get_pkt_length(),
                               self.identification, self.flagment0ffset, self.TTL,
                               self.protocol, self.headCheckSum, self.sourceBuff,
                               self.destinationBuff)
        return tempbuff
    def get_ip_checksum(self):
        checksum = 0
        IP_head = self._temp_pack()
        headlen = len(IP_head)
        i=0
        while i<headlen:
            temp = struct.unpack('!H',IP_head[i:i+2])[0]
            checksum = checksum+temp
            i = i+2
        checksum = (checksum>>16) + (checksum&0xffff)
        checksum = checksum+(checksum>>16)
        return 0xffff-checksum
    def pack(self):
        #insert the ip head checksums
        self.headCheckSum = self.get_ip_checksum()
        buff = self._temp_pack() + self.data
        return buff

class udp(ip_pkt):
    def __init__(self, sip, dip, sport, dport):
        ip_pkt.__init__(self, sip, dip)
        self.protocol = 0x11
        self.sport = sport
        self.dport = dport
        self.payload = ''
    def get_pkt_length(self):
        return 20 + 8 + len(self.payload)
    def pack(self):
        self.headCheckSum = self.get_ip_checksum()
        ipheader = self._temp_pack()
        udpheader = struct.pack('!HHHH', self.sport, self.dport, 8 + len(self.payload), 0)
        buff = ipheader + udpheader + self.payload
        return buff

class tcp(ip_pkt):
    def __init__(self, sip, dip, sport, dport):
        ip_pkt.__init__(self, sip, dip)
        self.protocol = 0x06
        self.sport = sport
        self.dport = dport
        self.sequenceNumber = 0xa07e21c6
        self.acknowlegementNumber = 0x00000000
        self.tcpOptionsBuff = ''
        self.tcpHeaderRealLength = 20
        self.tcpFlags = 0x02
        self.windowSize = 0x2000
        self.tcpCheckSum = 0x0000
        self.urgentPointer = 0x0000
        self.payload = ''
    def add_tcp_options(self, option_data):
        '''option_data   format:  020405b401010402 '''
        self.tcpOptionsBuff = binascii.a2b_hex(option_data)
        self.tcpHeaderRealLength = 20 + len(self.tcpOptionsBuff)
    def set_tcp_flags(self, flags):
        self.tcpFlags = flags
    def tcp_head_pack(self):
        # no check sum
        # tcp 6位保留位不处理，因此tcp头部长度字段和标志位字段直接拼接了
        buff = struct.pack('!HHLLBBHHH', self.sport, self.dport, self.sequenceNumber,
                           self.acknowlegementNumber, (self.tcpHeaderRealLength / 4) << 4, self.tcpFlags,
                           self.windowSize,
                           self.tcpCheckSum, self.urgentPointer)
        if self.tcpOptionsBuff:
            buff += self.tcpOptionsBuff
        return buff
    def get_tcp_check_sum(self):
        # fakehead buff (no pack reserved section)
        '''
        TCP 校验和的算法：
          ---------伪首部BUFF构造---------
          1、 SIP + DIP
          2、 IP头部中 reserved::protocol字段(0x0006，保留字段00在高位） + TCP报文长度(TCP头部+payload长度，2个字节方式pack)
          ---------以上两段buff拼接后得到伪首部的BUFF------------
          3、 TCP头部（含option字段）
          4、TCP payload （暂时还没有加到tcp头部校验中去）
          5、以上pack并拼接好的BUFF每16位相加，pyload不足16位的，左移8位。 得到的值  低16位 + 高位，最后取反
        '''
        fakeHead = self.sourceBuff + self.destinationBuff
        fakeHead += struct.pack('!HH', 0x06, (self.tcpHeaderRealLength + len(self.payload)))
        # fakehead END
        tcpHead = self.tcp_head_pack()
        if len(self.payload) % 2:
            tcpdata = self.payload + binascii.a2b_hex('00')
        else:
            tcpdata = self.payload
        data = fakeHead + tcpHead + tcpdata
        checksum = 0
        i = 0
        while i < len(data):
            temp = struct.unpack('!H', data[i:i + 2])[0]
            checksum = checksum + temp
            i = i + 2
        checksum = (checksum >> 16) + (checksum & 0xffff)
        checksum += (checksum >> 16)
        return 0xffff - checksum
    def get_pkt_length(self):
        # 得到IP头部的长度
        return 20 + len(self.tcp_head_pack()) + len(self.payload)
    def pack(self):
        self.headCheckSum = self.get_ip_checksum()
        ipheader = self._temp_pack()
        self.tcpCheckSum = self.get_tcp_check_sum()
        buff = ipheader + self.tcp_head_pack() + self.payload
        return buff

class buildPacket(tcp, udp):
    def __init__(self, eth_src, eth_dst, sip, dip, pro, sport, dport):
        tcp.__init__(self, sip, dip, sport, dport)
        udp.__init__(self, sip, dip, sport, dport)
        self.srcMac = eth_src
        self.dstMac = eth_dst
        self.ethType = 0x0800
        self.sip = sip
        self.dip = dip
        self.protocol = pro
        self.payload = ''
    def get_pkt_length(self):
        # 得到IP头部的长度
        if self.protocol == 6:
            return 20 + len(self.tcp_head_pack()) + len(self.payload)
        elif self.protocol == 17:
            return 20 + 8 + len(self.payload)
    def add_payload(self, data):
        #not include tcp or udp header
        self.payload = data
    def build_ethernet_head(self):
        srcMacBuff = binascii.a2b_hex(''.join(self.srcMac.split(':')))
        dstMacBuff = binascii.a2b_hex(''.join(self.dstMac.split(':')))
        ethTypeBuff = struct.pack('!H', self.ethType)
        return dstMacBuff + srcMacBuff + ethTypeBuff
    def pack(self):
        ethBuff = self.build_ethernet_head()
        if self.protocol == 17:
            print 'create a udp packet!'
            self.headCheckSum = self.get_ip_checksum()
            ipheader = self._temp_pack()
            udpheader = struct.pack('!HHHH', self.sport, self.dport, 8 + len(self.payload), 0)
            buff = ethBuff + ipheader + udpheader + self.payload
            return buff
        elif self.protocol == 6:
            print 'create a tcp packet!'
            self.headCheckSum = self.get_ip_checksum()
            ipheader = self._temp_pack()
            self.tcpCheckSum = self.get_tcp_check_sum()
            buff = ethBuff + ipheader + self.tcp_head_pack() + self.payload
            return buff
        else:
            buff = ''
            return
        return buff

def format_output(stream_buff = None, sizeof_eachline = 16):
    #sizef_each 为可选，定义每一行显示多少个字节的数据，一定要为8的倍数
    res = []
    temp_str = []
    for i, j in enumerate(stream_buff):
        #如果j为特殊字符或者超出ascii码，则用"."代替来显示
        if ord(j) < 32 or ord(j)>126:
            temp_str.append('.')
        else:
            temp_str.append(j)
        #每行输出16个字节的16进制和对应的ascii码值
        if i!=0 and i%sizeof_eachline == 0:
            res.append(' '*8)
            res.append(''.join(temp_str))
            temp_str = []
            res.append('\n')
        #每8个字节后，空3格
        elif i!=0 and i%8 == 0:
            res.append(' '*3)
        res.append(' ' + binascii.b2a_hex(j))
    #下面处理最后一行   space_count为16进制数据内容与buff内容显示隔开的空格数
    a, b = divmod(len(stream_buff)%sizeof_eachline, 8)
    if len(stream_buff)%sizeof_eachline == 0:
        space_count = 0
    elif not b:
        space_count = (sizeof_eachline - len(stream_buff) % sizeof_eachline + sizeof_eachline / 8 - a) * 3
    else:
        space_count = (sizeof_eachline - len(stream_buff) % sizeof_eachline + sizeof_eachline / 8 - a - 1) *3
    res.append(' '*(space_count+8) + ''.join(temp_str))
    return ''.join(res)


if __name__ == '__main__':
    a = buildPacket('00:00:00:00:00:01', '00:00:00:00:00:01', '1.1.1.1', '1.1.1.2', 17, 1000 ,64)
    b = buildPacket('00:00:00:00:00:01', '00:00:00:00:00:01', '1.1.1.1', '1.1.1.2', 6, 1000, 80)
    #print repr(a.build())
    #print repr(b.build())
    print format_output(a.pack())
    print format_output(b.pack())